#include "pin.H"
#include <unistd.h> // for pid
#include <stdlib.h>
#include <iostream>
#include <iomanip>
#include <fstream>
#include <map>

/* ===================================================================== */
/* Commandline Switches */
/* ===================================================================== */

KNOB<string> KnobOutputFile(KNOB_MODE_WRITEONCE,         "pintool",
                            "o", "hw1.out", "specify profile file name");

KNOB<BOOL>   KnobPid(KNOB_MODE_WRITEONCE,                "pintool",
                            "i", "0", "append pid to output");

KNOB<UINT64> KnobBranchLimit(KNOB_MODE_WRITEONCE,        "pintool",
                            "l", "0", "set limit of branches analyzed");

/* ===================================================================== */
/* Global Variables */
/* ===================================================================== */
UINT64 CountSeen = 0;
UINT64 IndirectCountSeen = 0;
UINT64 CountTaken = 0;
UINT64 CountCorrect = 0;
UINT64 IndirectCountCorrect = 0;
UINT64 CountReplaced = 0;
UINT64 IndirectCountReplaced = 0;
UINT64 ReturnCount= 0;
UINT64 DirectBranchCount = 0;
UINT64 DirectCallCount = 0;
UINT64 IndirectBranchCount= 0;
UINT64 IndirectCallCount= 0;
UINT64 InstructionCount= 0;

/* ===================================================================== */
/* Branch predictor                                                      */
/* ===================================================================== */
UINT64 mask = 0x03FF;
#define SIZE 1024

struct entry_one_bit
{
    bool valid;
    bool prediction;
    //UINT8 counter;
    UINT64 tag;
    UINT64 ReplaceCount;
    ADDRINT TargetAddr;
} BTB_one_bit[SIZE];

/* initialize the BTB */
VOID BTB_init()
{
    int i;

    for(i = 0; i < SIZE; i++)
    {
        BTB_one_bit[i].valid = false;
        BTB_one_bit[i].prediction = false;
        //BTB_one_bit[i].counter = 0;
        BTB_one_bit[i].tag = 0;
        BTB_one_bit[i].ReplaceCount = 0;
        BTB_one_bit[i].TargetAddr = 0;
    }
}

/* see if the given address is in the BTB */
bool BTB_lookup(ADDRINT ins_ptr)
{
    UINT64 index;

    index = mask & ins_ptr;

    if(BTB_one_bit[index].valid)
        if(BTB_one_bit[index].tag == ins_ptr)
            return true;

    return false;
}

/* return the prediction for the given address */
bool BTB_prediction(ADDRINT ins_ptr)
{
    UINT64 index;

    index = mask & ins_ptr;
    
    return BTB_one_bit[index].prediction;
}

ADDRINT BTB_prediction_indirect(ADDRINT ins_ptr)
{
    UINT64 index;

    index = mask & ins_ptr;
    
    return BTB_one_bit[index].TargetAddr;
}

/* update the BTB entry with the last result */
VOID BTB_update(ADDRINT ins_ptr, bool taken)
{
    UINT64 index;

    index = mask & ins_ptr;

    BTB_one_bit[index].prediction = taken;
}

VOID BTB_update_indirect(ADDRINT ins_ptr, bool taken, ADDRINT tar_addr)
{
    UINT64 index;

    index = mask & ins_ptr;

    BTB_one_bit[index].prediction = true;
    BTB_one_bit[index].TargetAddr = tar_addr;

}

/* insert a new branch in the table */
VOID BTB_insert(ADDRINT ins_ptr)
{
    UINT64 index;

    index = mask & ins_ptr;

    if(BTB_one_bit[index].valid)
    {
        BTB_one_bit[index].ReplaceCount++;
        CountReplaced++;
    }

    BTB_one_bit[index].valid = true;
    BTB_one_bit[index].prediction = true;
    //BTB_one_bit[index].counter = 2;
    BTB_one_bit[index].tag = ins_ptr;
}

VOID BTB_insert_indirect(ADDRINT ins_ptr, ADDRINT tar_addr)
{
    UINT64 index;

    index = mask & ins_ptr;

    if(BTB_one_bit[index].valid)
    {
        BTB_one_bit[index].ReplaceCount++;
        IndirectCountReplaced++;
    }

    BTB_one_bit[index].valid = true;
    BTB_one_bit[index].prediction = true;
    //BTB_one_bit[index].counter = 2;
    BTB_one_bit[index].tag = ins_ptr;
    BTB_one_bit[index].TargetAddr = tar_addr;
}
/* ===================================================================== */


/* ===================================================================== */

static INT32 Usage()
{
    cerr << "This pin tool collects a profile of jump/branch/call instructions for an application\n";

    cerr << KNOB_BASE::StringKnobSummary();

    cerr << endl;
    return -1;
}

/* ===================================================================== */

VOID write_results_indirect(bool limit_reached)
{
    int i;

    string output_file = KnobOutputFile.Value();
    if(KnobPid) output_file += "." + decstr(getpid());
    
    std::ofstream out(output_file.c_str());

    if(limit_reached)
        out << "Reason: limit reached\n";
    else
        out << "Reason: fini\n";
    out << "Indirect Count Seen: " << IndirectCountSeen << endl;
    out << "Indirect Count Correct: " << IndirectCountCorrect << endl;
    out << "Indirect Count Replaced: " << IndirectCountReplaced << endl;
    for(i = 0; i < SIZE; i++)
    {
        out << "BTB entry: " << i << ";" << BTB_one_bit[i].valid << ";" << BTB_one_bit[i].ReplaceCount << endl;
    }
    out.close();
}



VOID write_results(bool limit_reached)
{
    int i;

    string output_file = KnobOutputFile.Value();
    if(KnobPid) output_file += "." + decstr(getpid());
		//cerr << "write_result!" << endl;
    
    std::ofstream out("dump_file");

    /*if(limit_reached)
        out << "Reason: limit reached\n";
    else
        out << "Reason: fini\n";*/
    out << "Total Instruction Count: " << InstructionCount << endl;
    out << "Count Seen: " << CountSeen << endl;
    out << "Count Taken: " << CountTaken << endl;
    out << "Count Correct: " << CountCorrect << endl;
    out << "Count Replaced: " << CountReplaced << endl;
    for(i = 0; i < SIZE; i++)
    {
        //out << "BTB entry: " << i << ";" << BTB_one_bit[i].valid << ";" << BTB_one_bit[i].ReplaceCount << endl;
    }
    //write_results_indirect(true);
    //out << "Indirect Branch Stats\n";
    out << "Indirect Count Seen: " << IndirectCountSeen << endl;
    out << "Indirect Count Correct: " << IndirectCountCorrect << endl;
    out << "Indirect Count Replaced: " << IndirectCountReplaced << endl;
    out << "Return DirectCall DirectBranch IndirectCall IndirectBranch" << endl;
    out << ReturnCount << ";" << DirectCallCount << ";" << DirectBranchCount << ";" << IndirectCallCount << ";" << IndirectBranchCount << ";" << endl;
    out.close();
    //std::ofstream out2("hw1.out2");
    
}




/* ===================================================================== */

VOID br_predict(ADDRINT ins_ptr, INT32 taken)
{
    CountSeen++;
    if (taken)
        CountTaken++;

    if(BTB_lookup(ins_ptr)) 
    {
        if(BTB_prediction(ins_ptr) == taken) CountCorrect++;
        BTB_update(ins_ptr, taken);
    }
    else
    {
        if(!taken) CountCorrect++;
        else BTB_insert(ins_ptr);
    }

    if(CountSeen == KnobBranchLimit.Value())
    {
        write_results(true);
        exit(0);
    }
} 

VOID br_predict_indirect(ADDRINT ins_ptr, INT32 taken, ADDRINT target_address)
{
    IndirectCountSeen++;

    if(BTB_lookup(ins_ptr)) 
    {
        if(BTB_prediction_indirect(ins_ptr) == target_address) IndirectCountCorrect++;
        BTB_update_indirect(ins_ptr, taken, target_address);
    }
    else
    {
        BTB_insert_indirect(ins_ptr, target_address);
    }

    if(IndirectCountSeen == KnobBranchLimit.Value())
    {
        write_results_indirect(true);
        exit(0);
    }
} 

//  IARG_INST_PTR   
// ADDRINT ins_ptr

/* ===================================================================== */
VOID IncReturn(){
  ReturnCount ++;
}

VOID IncDirectCall(){
  DirectCallCount ++;
}

VOID IncDirectBranch(){
  DirectBranchCount ++;
}

VOID IncIndirectCall(){
  IndirectCallCount ++;
}

VOID IncIndirectBranch(){
  IndirectBranchCount ++;
}


VOID ins_count() {
  InstructionCount ++; 
}

VOID Instruction(INS ins, void *v)
{

// The subcases of direct branch and indirect branch are
// broken into "call" or "not call".  Call is for a subroutine
// These are left as subcases in case the programmer wants
// to extend the statistics to see how sub cases of branches behave
    INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) ins_count, IARG_END); 

    if( INS_IsRet(ins) )
    {
        INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) br_predict, 
            IARG_INST_PTR, IARG_BRANCH_TAKEN,  IARG_END);
        INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) IncReturn, IARG_END);
    }
    else if( INS_IsSyscall(ins) )
    {
        INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) br_predict, 
            IARG_INST_PTR, IARG_BRANCH_TAKEN,  IARG_END);
    }
    else if (INS_IsDirectBranchOrCall(ins))
    {
        if( INS_IsCall(ins) ) {
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) br_predict, 
                IARG_INST_PTR, IARG_BRANCH_TAKEN,  IARG_END);
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) IncDirectCall, IARG_END);
        }
        else {
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) br_predict, 
                IARG_INST_PTR, IARG_BRANCH_TAKEN,  IARG_END);
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) IncDirectBranch, IARG_END);
        }
    }
    else if( INS_IsIndirectBranchOrCall(ins) )
    {
        if( INS_IsCall(ins) ) {
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) br_predict_indirect, 
                IARG_INST_PTR, IARG_BRANCH_TAKEN,IARG_BRANCH_TARGET_ADDR,  IARG_END);
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) IncIndirectCall, IARG_END);
    }
        else {
          INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) IncIndirectBranch, IARG_END);
          if (INS_HasFallThrough(ins))
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) br_predict, 
                IARG_INST_PTR, IARG_BRANCH_TAKEN,  IARG_END);
          else
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) br_predict_indirect, 
                IARG_INST_PTR, IARG_BRANCH_TAKEN,IARG_BRANCH_TARGET_ADDR,  IARG_END);
        }
    }

}

/* ===================================================================== */

#define OUT(n, a, b) out << n << " " << a << setw(16) << CountSeen. b  << " " << setw(16) << CountTaken. b << endl
int Fini_count = 0;

VOID Fini(int n, void *v)
{
		cerr << "write_result!" << Fini_count << endl;
		Fini_count ++;
    write_results(false);
}


/* ===================================================================== */


/* ===================================================================== */

int main(int argc, char *argv[])
{
    
    if( PIN_Init(argc,argv) )
    {
        return Usage();
    }

    BTB_init();
        
    INS_AddInstrumentFunction(Instruction, 0);
    PIN_AddFiniFunction(Fini, 0);

    // Never returns

    PIN_StartProgram();
    
    return 0;
}

/* ===================================================================== */
/* eof */
/* ===================================================================== */
